package org.jsoup.helper;

import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.parser.Parser;
import org.jsoup.parser.TokenQueue;


import java.io.*;
import java.net.HttpURLConnection;
import java.net.InetSocketAddress;
import java.net.MalformedURLException;
import java.net.PasswordAuthentication;
import java.net.Proxy;
import java.net.URL;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.zip.GZIPInputStream;
import java.net.Authenticator;

/**
 * Implementation of {@link Connection}.
 * @see org.jsoup.Jsoup#connect(String) 
 */
public class HttpConnection  {
	public enum Method {
        GET, POST
    }
	
	private static HttpConnection.Request req;
    private static HttpConnection.Response res;
    
    private HttpConnection() {
        req = new Request();
        res = new Response();
    }
	
    public static HttpConnection connect(String url) {
    	HttpConnection con = new HttpConnection();
    	Validate.notEmpty(url, "Must supply a valid URL");
        try {
            req.url(new URL(url));
        } catch (MalformedURLException e) {
            throw new IllegalArgumentException("Malformed URL: " + url, e);
        }
        return con;
    }
    
    public HttpConnection userAgent(String userAgent) {
        Validate.notNull(userAgent);
        req.header("User-Agent", userAgent);
        return this;
    }
    
	public HttpConnection proxy(String hostname, int port, String username,
			String password) {
		
		Validate.notNull(hostname);
		Validate.notNull(port);
		
		res.setProxy(hostname, port, username, password);
		return null;
	}
	
    public HttpConnection timeout(int millis) {
        req.timeout(millis);
        return this;
    }
    
    public HttpConnection followRedirects(boolean followRedirects) {
        req.followRedirects(followRedirects);
        return this;
    }

    public HttpConnection referrer(String referrer) {
        Validate.notNull(referrer, "Referrer must not be null");
        req.header("Referer", referrer);
        return this;
    }

    public HttpConnection method(Method method) {
        req.method(method);
        return this;
    }

    public HttpConnection ignoreHttpErrors(boolean ignoreHttpErrors) {
		req.ignoreHttpErrors(ignoreHttpErrors);
		return this;
	}

    public HttpConnection ignoreContentType(boolean ignoreContentType) {
        req.ignoreContentType(ignoreContentType);
        return this;
    }

    public HttpConnection data(String key, String value) {
        req.data(KeyVal.create(key, value));
        return this;
    }

    public HttpConnection data(Map<String, String> data) {
        Validate.notNull(data, "Data map must not be null");
        for (Map.Entry<String, String> entry : data.entrySet()) {
            req.data(KeyVal.create(entry.getKey(), entry.getValue()));
        }
        return this;
    }

    public HttpConnection data(String... keyvals) {
        Validate.notNull(keyvals, "Data key value pairs must not be null");
        Validate.isTrue(keyvals.length %2 == 0, "Must supply an even number of key value pairs");
        for (int i = 0; i < keyvals.length; i += 2) {
            String key = keyvals[i];
            String value = keyvals[i+1];
            Validate.notEmpty(key, "Data key must not be empty");
            Validate.notNull(value, "Data value must not be null");
            req.data(KeyVal.create(key, value));
        }
        return this;
    }

    public HttpConnection header(String name, String value) {
        req.header(name, value);
        return this;
    }
    
    public Document get() throws IOException {
        req.method(Method.GET);
        execute();
        return res.parse();
    }

    public Document post() throws IOException {
        req.method(Method.POST);
        execute();
        return res.parse();
    }

    public HttpConnection.Response execute() throws IOException {
        res = Response.execute(req);
        return res;
    }

    public HttpConnection.Request request() {
        return req;
    }
    
	public HttpConnection request(HttpConnection.Request request) {
        req = request;
        return this;
    }
	
    public HttpConnection.Response response() {
        return res;
    }
    
    public HttpConnection response(HttpConnection.Response response) {
        res = response;
        return this;
    }

    @SuppressWarnings({"unchecked", "rawtypes"})
    private static abstract class Base<T extends HttpConnection.Base> {
        URL url;
        Method method;
        Map<String, String> headers;
        Map<String, String> cookies;

        private Base() {
            headers = new LinkedHashMap<String, String>();
            cookies = new LinkedHashMap<String, String>();
        }
        
        public URL url() {
            return url;
        }
        
        public T url(URL url) {
            Validate.notNull(url, "URL must not be null");
            this.url = url;
            return (T) this;
        }
        
        public Method method() {
            return method;
        }

        public T method(Method method) {
            Validate.notNull(method, "Method must not be null");
            this.method = method;
            return (T) this;
        }

        public String header(String name) {
            Validate.notNull(name, "Header name must not be null");
            return getHeaderCaseInsensitive(name);
        }

        public T header(String name, String value) {
            Validate.notEmpty(name, "Header name must not be empty");
            Validate.notNull(value, "Header value must not be null");
            removeHeader(name); // ensures we don't get an "accept-encoding" and a "Accept-Encoding"
            headers.put(name, value);
            return (T) this;
        }

        public boolean hasHeader(String name) {
            Validate.notEmpty(name, "Header name must not be empty");
            return getHeaderCaseInsensitive(name) != null;
        }

        public T removeHeader(String name) {
            Validate.notEmpty(name, "Header name must not be empty");
            Map.Entry<String, String> entry = scanHeaders(name); // remove is case insensitive too
            if (entry != null)
                headers.remove(entry.getKey()); // ensures correct case
            return (T) this;
        }
        
        public Map<String, String> headers() {
            return headers;
        }
        
        private String getHeaderCaseInsensitive(String name) {
            Validate.notNull(name, "Header name must not be null");
            // quick evals for common case of title case, lower case, then scan for mixed
            String value = headers.get(name);
            if (value == null)
                value = headers.get(name.toLowerCase());
            if (value == null) {
                Map.Entry<String, String> entry = scanHeaders(name);
                if (entry != null)
                    value = entry.getValue();
            }
            return value;
        }
        
        private Map.Entry<String, String> scanHeaders(String name) {
            String lc = name.toLowerCase();
            for (Map.Entry<String, String> entry : headers.entrySet()) {
                if (entry.getKey().toLowerCase().equals(lc))
                    return entry;
            }
            return null;
        }
        
        public String cookie(String name) {
            Validate.notNull(name, "Cookie name must not be null");
            return cookies.get(name);
        }
        
        public T cookie(String name, String value) {
            Validate.notEmpty(name, "Cookie name must not be empty");
            Validate.notNull(value, "Cookie value must not be null");
            cookies.put(name, value);
            return (T) this;
        }

        public boolean hasCookie(String name) {
            Validate.notEmpty("Cookie name must not be empty");
            return cookies.containsKey(name);
        }

        public T removeCookie(String name) {
            Validate.notEmpty("Cookie name must not be empty");
            cookies.remove(name);
            return (T) this;
        }

        public Map<String, String> cookies() {
            return cookies;
        }
    }
    
    public class Request extends Base<HttpConnection.Request> {
        private int timeoutMilliseconds;
        private boolean followRedirects;
        private Collection<HttpConnection.KeyVal> data;
        private boolean ignoreHttpErrors = false;
        private boolean ignoreContentType = false;

      	private Request() {
            timeoutMilliseconds = 3000;
            followRedirects = true;
            data = new ArrayList<HttpConnection.KeyVal>();
            method = HttpConnection.Method.GET;
            headers.put("Accept-Encoding", "gzip");
        }

        public int timeout() {
            return timeoutMilliseconds;
        }
        
        public Request timeout(int millis) {
            Validate.isTrue(millis >= 0, "Timeout milliseconds must be 0 (infinite) or greater");
            timeoutMilliseconds = millis;
            return this;
        }

        public boolean followRedirects() {
            return followRedirects;
        }

        public HttpConnection.Request followRedirects(boolean followRedirects) {
            this.followRedirects = followRedirects;
            return this;
        }

        public boolean ignoreHttpErrors() {
            return ignoreHttpErrors;
        }

        public void ignoreHttpErrors(boolean ignoreHttpErrors) {
            this.ignoreHttpErrors = ignoreHttpErrors;
        }

        public boolean ignoreContentType() {
            return ignoreContentType;
        }

        public void ignoreContentType(boolean ignoreContentType) {
            this.ignoreContentType = ignoreContentType;
        }

        public Request data(HttpConnection.KeyVal keyval) {
            Validate.notNull(keyval, "Key val must not be null");
            data.add(keyval);
            return this;
        }

        public Collection<HttpConnection.KeyVal> data() {
            return data;
        }
    }

    public static class Response extends Base<HttpConnection.Response> {
        private static final int MAX_REDIRECTS = 20;
        private int statusCode;
        private String statusMessage;
        private ByteBuffer byteData;
        private String charset;
        private String contentType;
        private boolean executed = false;
        private int numRedirects = 0;
        private HttpConnection.Request req;
        private static boolean isProxySet = false;
        private static Proxy proxy = null;
        
        Response() {
            super();
        }
        
        private Response(Response previousResponse) throws IOException {
            super();
            if (previousResponse != null) {
                numRedirects = previousResponse.numRedirects + 1;
                if (numRedirects >= MAX_REDIRECTS)
                    throw new IOException(String.format("Too many redirects occurred trying to load URL %s", previousResponse.url()));
            }
        }

        static Response execute(HttpConnection.Request req) throws IOException {
            return execute(req, null);
        }

        static Response execute(HttpConnection.Request req, Response previousResponse) throws IOException {
            Validate.notNull(req, "Request must not be null");
            String protocol = req.url().getProtocol();
            Validate
                .isTrue(protocol.equals("http") || protocol.equals("https"), "Only http & https protocols supported");

            // set up the request for execution
            if (req.method() == HttpConnection.Method.GET && req.data().size() > 0)
                serialiseRequestUrl(req); // appends query string
            HttpURLConnection conn = createConnection(req);
            
            conn.connect();
            
            if (req.method() == HttpConnection.Method.POST)
                writePost(req.data(), conn.getOutputStream());          
            
            int status = conn.getResponseCode();
            boolean needsRedirect = false;
            if (status != HttpURLConnection.HTTP_OK) {
                if (status == HttpURLConnection.HTTP_MOVED_TEMP || status == HttpURLConnection.HTTP_MOVED_PERM || status == HttpURLConnection.HTTP_SEE_OTHER)
                    needsRedirect = true;
                else if (!req.ignoreHttpErrors())
                    throw new IOException(status + " error loading URL " + req.url().toString());
            }
            Response res = new Response(previousResponse);
            res.setupFromConnection(conn, previousResponse);
            if (needsRedirect && req.followRedirects()) {
                req.url(new URL(req.url(), res.header("Location")));
                
                return execute(req, res);
            }
            res.req = req;

            InputStream bodyStream = null;
            InputStream dataStream = null;
            try {
                dataStream = conn.getErrorStream() != null ? conn.getErrorStream() : conn.getInputStream();
            	bodyStream = res.hasHeader("Content-Encoding") && res.header("Content-Encoding").equalsIgnoreCase("gzip") ?
                        new BufferedInputStream(new GZIPInputStream(dataStream)) :
                        new BufferedInputStream(dataStream);
                
                res.byteData = readToByteBuffer(bodyStream);
                res.charset = getCharsetFromContentType(res.contentType); // may be null, readInputStream deals with it
            } finally {
                if (bodyStream != null) bodyStream.close();
                if (dataStream != null) dataStream.close();
            }

            res.executed = true;
            return res;
        }
        
        /**
         * Parse out a charset from a content type header.
         * @param contentType e.g. "text/html; charset=EUC-JP"
         * @return "EUC-JP", or null if not found. Charset is trimmed and uppercased.
         */
        static String getCharsetFromContentType(String contentType) {
            if (contentType == null) return null;
            Matcher m = Pattern.compile("(?i)\\bcharset=\\s*\"?([^\\s;\"]*)").matcher(contentType);
            if (m.find()) {
                return m.group(1).trim().toUpperCase();
            }
            return null;
        }
        
        static ByteBuffer readToByteBuffer(InputStream inStream) throws IOException {
            byte[] buffer = new byte[0x20000];
            ByteArrayOutputStream outStream = new ByteArrayOutputStream(0x20000);
            int read;
            while(true) {
                read  = inStream.read(buffer);
                if (read == -1) break;
                outStream.write(buffer, 0, read);
            }
            ByteBuffer byteData = ByteBuffer.wrap(outStream.toByteArray());
            return byteData;
        }
        
        public void setProxy(String proxyHostname, int port, String username, String password) {
        	ProxyAuthenticator pa = new ProxyAuthenticator(username, password);
        	Authenticator.setDefault(pa);
        	
            proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(proxyHostname, port));
            isProxySet = true;
        }
        
        public int statusCode() {
            return statusCode;
        }

        public String statusMessage() {
            return statusMessage;
        }

        public String charset() {
            return charset;
        }

        public String contentType() {
            return contentType;
        }

        public Document parse() throws IOException {
            Validate.isTrue(executed, "Request must be executed (with .execute(), .get(), or .post() before parsing response");
            if (!req.ignoreContentType() && (contentType == null || !(contentType.startsWith("text/") || contentType.startsWith("application/xml") || contentType.startsWith("application/xhtml+xml"))))
                throw new IOException(String.format("Unhandled content type \"%s\" on URL %s. Must be text/*, application/xml, or application/xhtml+xml",
                    contentType, url.toString()));
            
            //can I rip out HTML for specific pages?
            
            
            Document doc = parseByteData(byteData, charset, url.toExternalForm());
            byteData.rewind();
            charset = doc.outputSettings().charset().name(); // update charset from meta-equiv, possibly
            return doc;
        }
        
        // reads bytes first into a buffer, then decodes with the appropriate charset. done this way to support
        // switching the chartset midstream when a meta http-equiv tag defines the charset.
        static Document parseByteData(ByteBuffer byteData, String charsetName, String baseUri) {
            String docData;
            Document doc = null;
            if (charsetName == null) { // determine from meta. safe parse as UTF-8
                // look for <meta http-equiv="Content-Type" content="text/html;charset=gb2312"> or HTML5 <meta charset="gb2312">
                docData = Charset.forName("UTF-8").decode(byteData).toString();
                doc = Parser.parse(docData, baseUri);
                Element meta = doc.select("meta[http-equiv=content-type], meta[charset]").first();
                if (meta != null) { // if not found, will keep utf-8 as best attempt
                    String foundCharset = meta.hasAttr("http-equiv") ? getCharsetFromContentType(meta.attr("content")) : meta.attr("charset");
                    if (foundCharset != null && foundCharset.length() != 0 && !foundCharset.equals("UTF-8")) { // need to re-decode
                        charsetName = foundCharset;
                        byteData.rewind();
                        docData = Charset.forName(foundCharset).decode(byteData).toString();
                        doc = null;
                    }
                }
            } else { // specified by content type header (or by user on file load)
                Validate.notEmpty(charsetName, "Must set charset arg to character set of file to parse. Set to null to attempt to detect from HTML");
                docData = Charset.forName(charsetName).decode(byteData).toString();
            }
            if (doc == null) {
                doc = Parser.parse(docData, baseUri);
                doc.outputSettings().charset(charsetName);
            }
            return doc;
        }

        public String body() {
            Validate.isTrue(executed, "Request must be executed (with .execute(), .get(), or .post() before getting response body");
            // charset gets set from header on execute, and from meta-equiv on parse. parse may not have happened yet
            String body;
            if (charset == null)
                body = Charset.forName("UTF-8").decode(byteData).toString();
            else
                body = Charset.forName(charset).decode(byteData).toString();
            byteData.rewind();
            return body;
        }

        public byte[] bodyAsBytes() {
            Validate.isTrue(executed, "Request must be executed (with .execute(), .get(), or .post() before getting response body");
            return byteData.array();
        }
        
        // set up connection defaults, and details from request
        private static HttpURLConnection createConnection(HttpConnection.Request req) throws IOException {
        	HttpURLConnection conn = (isProxySet) 
        		? (HttpURLConnection) req.url().openConnection(proxy) 
        		: (HttpURLConnection) req.url().openConnection();
        	
            conn.setRequestMethod(req.method().name());
            conn.setInstanceFollowRedirects(false); // don't rely on native redirection support
            conn.setConnectTimeout(req.timeout());
            conn.setReadTimeout(req.timeout());
            if (req.method() == Method.POST)
                conn.setDoOutput(true);
            
            for (Map.Entry<String, String> header : req.headers().entrySet()) {
                conn.addRequestProperty(header.getKey(), header.getValue());
            }
            return conn;
        }
        
        // set up url, method, header, cookies
        private void setupFromConnection(HttpURLConnection conn, HttpConnection.Response previousResponse) throws IOException {
            method = HttpConnection.Method.valueOf(conn.getRequestMethod());
            url = conn.getURL();
            statusCode = conn.getResponseCode();
            statusMessage = conn.getResponseMessage();
            contentType = conn.getContentType();

            Map<String, List<String>> resHeaders = conn.getHeaderFields();
            processResponseHeaders(resHeaders);
            
        }

        void processResponseHeaders(Map<String, List<String>> resHeaders) {
            for (Map.Entry<String, List<String>> entry : resHeaders.entrySet()) {
                String name = entry.getKey();
                if (name == null)
                    continue; // http/1.1 line

                List<String> values = entry.getValue();
                if (name.equalsIgnoreCase("Set-Cookie")) {
                    for (String value : values) {
                        if (value == null)
                            continue;
                        TokenQueue cd = new TokenQueue(value);
                        String cookieName = cd.chompTo("=").trim();
                        String cookieVal = cd.consumeTo(";").trim();
                        if (cookieVal == null)
                            cookieVal = "";
                        // ignores path, date, domain, secure et al. req'd?
                        // name not blank, value not null
                        if (cookieName != null && cookieName.length() > 0)
                            cookie(cookieName, cookieVal);
                    }
                } else { // only take the first instance of each header
                    if (!values.isEmpty())
                        header(name, values.get(0));
                }
            }
        }

        private static void writePost(Collection<HttpConnection.KeyVal> data, OutputStream outputStream) throws IOException {
            OutputStreamWriter w = new OutputStreamWriter(outputStream, "UTF-8");
            boolean first = true;
            for (HttpConnection.KeyVal keyVal : data) {
                if (!first) 
                    w.append('&');
                else
                    first = false;
                
                w.write(URLEncoder.encode(keyVal.key(), "UTF-8"));
                w.write('=');
                w.write(URLEncoder.encode(keyVal.value(), "UTF-8"));
            }
            w.close();
        }
         
        // for get url reqs, serialise the data map into the url
        private static void serialiseRequestUrl(HttpConnection.Request req) throws IOException {
            URL in = req.url();
            StringBuilder url = new StringBuilder();
            boolean first = true;
            // reconstitute the query, ready for appends
            url
                .append(in.getProtocol())
                .append("://")
                .append(in.getAuthority()) // includes host, port
                .append(in.getPath())
                .append("?");
            if (in.getQuery() != null) {
                url.append(in.getQuery());
                first = false;
            }
            for (HttpConnection.KeyVal keyVal : req.data()) {
                if (!first)
                    url.append('&');
                else
                    first = false;
                url
                    .append(URLEncoder.encode(keyVal.key(), "UTF-8"))
                    .append('=')
                    .append(URLEncoder.encode(keyVal.value(), "UTF-8"));
            }
            req.url(new URL(url.toString()));
            req.data().clear(); // moved into url as get params
        }
    }
    
    public static class KeyVal {
        private String key;
        private String value;

        public static KeyVal create(String key, String value) {
            Validate.notEmpty(key, "Data key must not be empty");
            Validate.notNull(value, "Data value must not be null");
            return new KeyVal(key, value);
        }

        private KeyVal(String key, String value) {
            this.key = key;
            this.value = value;
        }

        public KeyVal key(String key) {
            Validate.notEmpty(key, "Data key must not be empty");
            this.key = key;
            return this;
        }

        public String key() {
            return key;
        }

        public KeyVal value(String value) {
            Validate.notNull(value, "Data value must not be null");
            this.value = value;
            return this;
        }

        public String value() {
            return value;
        }

        @Override
        public String toString() {
            return key + "=" + value;
        }      
    }
    
    public static class ProxyAuthenticator extends Authenticator {

        private String user, password;
        
        public ProxyAuthenticator(String user, String password) {
            this.user = user;
            this.password = password;
        }
        
        protected PasswordAuthentication getPasswordAuthentication() {
            return new PasswordAuthentication(user, password.toCharArray());
        }
    }
}

